import os
import sys

import bpy
from bpy.types import NodeTree, Node, NodeSocket

pwd = os.getcwd()
sys.path.append(pwd)

from vad_data import *


# blender add-on部分——tool shelf中的Dynamic display

class memoryDynamicPanel(bpy.types.Panel):
    bl_idname = 'memory.DynamicDisplay'
    bl_label = 'Dynamic Display'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        col = self.layout.column(align=True)
        col.label(text="Dynamic control:")
        row = col.row(align=True)
        row.operator("memory.prev", text="Prev")
        row.operator("memory.next", text="Next")

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))

#
#    控制前进，动态显示
#
class MEMORY_OT_NextButton(bpy.types.Operator):
    bl_idname = "memory.next"
    bl_label = "Next Memory"
    
    def execute(self, context):
        pid = "408"
        vad_node = VadData(pid)
        vadtree = vad_node.next()

        if vadtree:
            tree = bpy.context.scene.node_tree
            # clear nodes
            for node in tree.nodes:
                tree.nodes.remove(node)
            
            draw = DrawVads()
            draw.drawVadNode(vadtree, tree)
            
        return{'FINISHED'}

#
#   控制后退
#
class MEMORY_OT_PrevButton(bpy.types.Operator):
    bl_idname = "memory.prev"
    bl_label = "previous Memory"
    
    def execute(self, context):
        pid = "408"
        vad_node = VadData(pid)
        vadtree = vad_node.prev()

        if vadtree:
            tree = bpy.context.scene.node_tree
            # clear nodes
            for node in tree.nodes:
                tree.nodes.remove(node)
            
            draw = DrawVads()
            draw.drawVadNode(vadtree, tree)
        return{'FINISHED'}


'''
draw object of vad graph
'''
class DrawVads(object, metaclass=Singleton):

    def __init__(self):
        self.imageNum = 0
        self.vadGraphs = []
        # node color
        self.colors = { "red": (1.0, 0.411, 0.503), "yellow": (0.981, 1.0, 0.272), "green": (0.563, 1.0, 0.541),
           "blue": (0.241, 0.861, 1.0), "gray": (0.4, 0.4, 0.4), "white": (1.0, 1.0, 1.0), 
           "purple": (0.58, 0.579, 1.0)}

    def drawVadNode(self, vadtree, tree):
        # create input vad node
        for vad_label in vadtree["nodes"]:
            vad_node = tree.nodes.new(type='VadNodeType')
            vad_node.location = tuple(vadtree["nodes"][vad_label]["location"])
            vad_node.name = vad_label
            vad_node.label = vad_label
            vad_node.VAD = vadtree["nodes"][vad_label]["address"]
            vad_node.use_custom_color = True
            vad_node.color = self.colors[vadtree["nodes"][vad_label]["color"]]

        links = tree.links
        for e in vadtree["edges"]:
            source = tree.nodes[e["source"]]
            target = tree.nodes[e["target"]]
            links.new(source.outputs[0], target.inputs[0])



class memoryColorPanel(bpy.types.Panel):
    bl_idname = 'memory.ColorInfo'
    bl_label = 'Color Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        purple_ball = icon_collections["purple_ball"]["purple_ball"]
        gray_ball = icon_collections["gray_ball"]["gray_ball"]
        green_ball = icon_collections["green_ball"]["green_ball"]
        yellow_ball = icon_collections["yellow_ball"]["yellow_ball"]
        blue_ball = icon_collections["blue_ball"]["blue_ball"]
        red_ball = icon_collections["red_ball"]["red_ball"]
        white_ball = icon_collections["white_ball"]["white_ball"]

        col = self.layout.column(align=True)
        col.label(text=": Heaps", icon_value=purple_ball.icon_id)
        col.label(text=": DLLs", icon_value=gray_ball.icon_id)
        col.label(text=": Stacks", icon_value=green_ball.icon_id)
        col.label(text=": Mapped Files", icon_value=yellow_ball.icon_id)
        col.label(text=": root node", icon_value=blue_ball.icon_id)
        col.label(text=": new or changed nodes", icon_value=red_ball.icon_id)
        col.label(text=": normal VAD nodes", icon_value=white_ball.icon_id)

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


class memoryInfoPanel(bpy.types.Panel):
    bl_idname = 'memory.VADInfo'
    bl_label = 'VAD Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Memory'

    def draw(self, context):
        pid="408"
        vad_node = VadData(pid)
        vadtree = vad_node.current()
        if context.active_node:
            node_info = vadtree["nodes"][context.active_node.name]
            col = self.layout.column(align=True)
            col.label(text="pid: "+str(node_info["pid"]))
            col.label(text="vad: "+hex(node_info["vad"]))
            col.label(text="start: "+hex(node_info["start"]))
            col.label(text="end: "+hex(node_info["end"]))
            col.label(text="VadTag: "+str(node_info["VadTag"]))
            col.label(text="VadType: "+str(node_info["VadType"]))
            col.label(icon="IPO_SINE")
            col.label(text="flags: "+str(node_info["flags"]))
            col.label(text="protection: "+str(node_info["protection"]))
            col.label(icon="IPO_QUAD")
            col.label(text="ControlArea: "+hex(node_info["ControlArea"]))
            col.label(text="segment: "+hex(node_info["segment"]))
            col.label(icon="IPO_CUBIC")
            col.label(text="NumberOfSectionReferences: "+str(node_info["NumberOfSectionReferences"]))
            col.label(text="NumberOfPfnReferences: "+str(node_info["NumberOfPfnReferences"]))
            col.label(text="NumberOfMappedViews: "+str(node_info["NumberOfMappedViews"]))
            col.label(text="NumberOfUserReferences: "+str(node_info["NumberOfUserReferences"]))
            col.label(icon="IPO_QUART")
            col.label(text="ControlFlags: "+str(node_info["ControlFlags"]))
            col.label(text="FileObject: "+hex(node_info["FileObject"]))
            col.label(text="FileName: "+str(node_info["FileName"]))
            col.label(icon="IPO_QUINT")
            col.label(text="FirstprototypePTE: "+hex(node_info["FirstprototypePTE"]))
            col.label(text="LastcontiguousPTE: "+hex(node_info["LastcontiguousPTE"]))
            col.label(text="Flags2: "+str(node_info["Flags2"]))
        

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


class VadCustomTree(NodeTree):
    '''A vad node tree type that will show up in the node editor header'''
    bl_idname = 'VadTreeType'
    bl_label = 'Vad Node Tree'
    bl_icon = 'NODETREE'

# Vad socket type
class VadCustomSocket(NodeSocket):
    '''Vad node socket type'''
    bl_idname = 'VadSocketType'
    bl_label = 'Vad Node Socket'

    def draw(self, context, layout, node, text):
        layout.label(text)

    def draw_color(self, context, node):
        return (1.0, 0.4, 0.216, 0.5)


class VadCustomTreeNode:
    @classmethod
    def poll(cls, ntree):
        return ntree.bl_idname == 'VadTreeType'


class OBJECT_OT_VADButton(bpy.types.Operator):
    bl_idname = "vad.address"
    bl_label = "Vad Address"
    
    def execute(self, context):
        return{'FINISHED'}


class VadCustomNode(Node, VadCustomTreeNode):
    '''A custom node'''
    bl_idname = 'VadNodeType'
    bl_label = 'Vad Node'
    bl_icon = 'SOUND'

    VAD = bpy.props.StringProperty(default="0x70238a30 - 0x7ffd2fff")

    def init(self, context):
        self.inputs.new('VadSocketType', " ")

        self.outputs.new('NodeSocketColor', " ")
    
    def copy(self, node):
        pass
    
    def free(self):
        pass
    
    def draw_buttons(self, context, layout):
        tmp = self.VAD
        layout.label(text=str(tmp))
    
    def draw_label(self):
        return "vad node"

# import nodeitems_utils
# from nodeitems_utils import NodeCategory, NodeItem

# class VadNodeCategory(NodeCategory):
#     @classmethod
#     def poll(cls, context):
#         return context.space_data.tree_type == 'VadTreeType'

# node_categories = [
#     # identifier, label, items list
#     VadNodeCategory("SOMENODES", "Some Nodes", items=[
#         # our basic node
#         NodeItem("VadNodeType"),
#         ]),
#     ]

# icon collection
icon_collections = {}

def register():
    import bpy.utils.previews
    icon = bpy.utils.previews.new()
    icons_dir = "../../icons"
    for ball in os.listdir(icons_dir):
        icon.load(ball[:-4], os.path.join(icons_dir, ball), 'IMAGE')
        icon_collections[ball[:-4]] = icon
    bpy.utils.register_class(memoryColorPanel)
    bpy.utils.register_class(memoryDynamicPanel)
    bpy.utils.register_class(MEMORY_OT_NextButton)
    bpy.utils.register_class(MEMORY_OT_PrevButton)
    bpy.utils.register_class(memoryInfoPanel)
    bpy.utils.register_class(OBJECT_OT_VADButton)
    bpy.utils.register_class(VadCustomTree)
    bpy.utils.register_class(VadCustomSocket)
    bpy.utils.register_class(VadCustomNode)

    # nodeitems_utils.register_node_categories("CUSTOM_NODES", node_categories)


def unregister():
    # nodeitems_utils.unregister_node_categories("CUSTOM_NODES")
    bpy.utils.unregister_class(OBJECT_OT_VADButton)
    bpy.utils.unregister_class(VadCustomTree)
    bpy.utils.unregister_class(VadCustomSocket)
    bpy.utils.unregister_class(VadCustomNode)
    bpy.utils.unregister_class(memoryDynamicPanel)
    bpy.utils.unregister_class(memoryInfoPanel)
    bpy.utils.unregister_class(MEMORY_OT_NextButton)
    bpy.utils.unregister_class(MEMORY_OT_PrevButton)
    for icon in icon_collections.values():
        bpy.utils.previews.remove(icon)
    icon_collections.clear()
    bpy.utils.unregister_class(memoryColorPanel)


if __name__ == "__main__":
    pid = "408"
    try:
        register()
    except Exception as e:
        print(e)
    ### blender operation ###
    bpy.context.scene.use_nodes = True
    tree = bpy.context.scene.node_tree

    # clear default nodes
    for node in tree.nodes:
        tree.nodes.remove(node)

    vad_node = VadData(pid)
    vadtree = vad_node.current()

    draw = DrawVads()
    draw.drawVadNode(vadtree, tree)

    
